import pickle
import torch
from chat_bot.models.NER.model.LSTM import NER_LSTM
from chat_bot.models.NER.model.LSTM_CRF import NER_LSTM_CRF
from chat_bot.models.NER.model.BERT_LSTM_CRF import NER_BERT_LSTM_CRF
from chat_bot.models.NER import config


class NERPredictor(object):
    def __init__(self):
        with open(config.pickle_path, 'rb') as f:
            self.word2id = pickle.load(f)
            self.id2word = pickle.load(f)
            self.tag2id = pickle.load(f)
            self.id2tag = pickle.load(f)

        self.models = {'NER_LSTM': NER_LSTM, 'NER_LSTM_CRF': NER_LSTM_CRF, 'NER_BERT_LSTM_CRF': NER_BERT_LSTM_CRF}

        self.model = self.models[config.model_name](
            embedding_dim=config.embedding_dim,
            hidden_dim=config.hidden_dim,
            dropout=config.dropout,
            word2id=self.word2id,
            tag2id=self.tag2id
        )

        self.model.load_state_dict(
            torch.load(
                f'{config.root_path}/output/{str(config.model_name).lower()}.pkl'
            )
        )

    def predict(self, text):
        input_vec = [self.word2id.get(i, 0) for i in text]
        sentences = torch.tensor(input_vec).view(1, -1)
        pre = self.model(sentences)
        pre = [self.id2tag[i] for i in pre[0]]
        entities = {}
        entity = ''
        entity_type = 'O'
        for i in range(len(pre)):
            if pre[i].startswith('B_'):
                if entity_type != 'O':
                    entities.setdefault(entity_type, []).append(entity)

                entity = text[i]
                entity_type = pre[i][2:]
            elif pre[i].startswith('I_') and entity_type == pre[i][2:]:
                entity += text[i]
            else:
                if entity_type != 'O':
                    entities.setdefault(entity_type, []).append(entity)
                entity = ''
                entity_type = 'O'
        else:
            if entity_type != 'O':
                entities.setdefault(entity_type, []).append(entity)

        return entities


if __name__ == '__main__':
    p = NERPredictor()
    while True:
        input_str = input("请输入文本: ")
        # output = predict(model, input_str)
        output = p.predict(input_str)

        # output = output[1]
        print(output)
